#!/usr/bin/env python 
# -*- coding:utf-8 -*-

from lib.utils.settings import *
from lib.utils.check import *
from lib.utils.printer import *
from lib.utils.loader import *
from os import sep
from time import strftime
import json
from lib.handler.report import *

execmod = []
resultJson = {}
#命令信息
def help_set():
    info('set命令参数:')
    plus("参数:url[*]")
    plus("参数:headers,设置头信息")
    plus("参数:auth,设置认证信息")
    plus("参数:cookie,设置cookie信息")
    plus("参数:agent,设置浏览器版本信息")
    plus("参数:redirect,关闭重定向")
    plus("参数:method,默认get提交方式")
    plus("参数:data,post提交方式时设置data数据")
    plus("参数:timeout,设置响应超时时间")
def info_plugins():
    info('插件命令和描述:')
    for k,v in plugs_info.items():
        s = "{}: {}".format(k,v)
        plus(s)
def info_plugins_simple():
    info('插件命令简写:')
    for k,v in plugs.items():
        s = "{}: {}".format(k,v)
        plus(s)
def info_attacks():
    info("attacks插件命令和描述:")
    for k,v in attacks_info.items():
        s = "{}: {}".format(k,v)
        plus(s)
def info_audit():
    info("audit插件命令和描述:")
    for k,v in audit_info.items():
        s = "{}: {}".format(k,v)
        plus(s)
def info_brute():
    info("brute插件命令和描述:")
    for k,v in brute_info.items():
        s = "{}: {}".format(k,v)
        plus(s)
def info_disclosure():
    info("disclosure插件命令和描述:")
    for k,v in disclosure_info.items():
        s = "{}: {}".format(k,v)
        plus(s)
#扫描配置
def set_url(url):
    ARGV['url'] = url
    info("url设置成功")

def set_headers(headers:str):
    headers = Cheaders(headers)
    if headers:
        warn("格式错误,例:k1:v1,k2:v2")
    else:
        ARGV['headers'] = headers
        info("头信息设置成功.")

def set_headers_host(host:str):
    ARGV['headers'].update({'Host':host})
    info("头信息[Host]设置成功")

def set_headers_referer(referer:str):
    ARGV['headers'].update({'Referer':referer})
    info("头信息[referer]设置成功")

##post提交方式需设置
def set_data(data:str):
    ARGV['data'] = data
    info("data数据设置成功")

def set_method(method:str):
    ARGV['method'] = method
    info("method方法设置成功")

def set_auth(auth:str):
    ARGV['auth'] = Cauth(auth)
    info("auth认证信息设置成功")

def set_agent(agent:str):
    ARGV['agent'] = agent
    info("agent浏览器版本信息设置成功!")

def set_random_agent(agent:str):
    ARGV['agent'] = ragent()
    info("agent浏览器版本信息设置成功!")

def set_cookie(cookie:str):
    ARGV['cookie'] = cookie
    info("cookie信息设置成功!")

def set_proxy(proxy:str):
    ARGV['proxy'] = ragent()
    info("代理设置成功!")

def set_proxy_auth(pauth:str):
    ARGV['pauth'] = Cauth(pauth)
    info("代理认证设置成功!")

def set_timeout(timeout:str):
    ARGV['timeout'] = float(timeout)
    info("超时时间设置成功!")

def set_redirect_false(redirect:bool):
    ARGV['redirect'] = False
    info("关闭重定向!")

#调用扫描模块
def startup_plugins(plugins:list):
    global resultJson
    for p in plugins:
        with open(p,"r",encoding="utf-8") as f:
            mod = load_string_to_module(f.read())
            resultJson.update( mod.run(ARGV,ARGV['url'],ARGV['data']))
def startup_spec_attacks(attack:str):
    if attack in attacks_info.keys():
        plugins = spec_attacks_plugins(attack)
        startup_plugins(plugins)
    else:
        warn("模块不存在!")

def startup_full_attacks():
    global resultJson
    if not REPORT['startTime']:
        REPORT['startTime'] = strftime("%Y/%m/%d at %H:%M:%S")
    execmod.append("attacks")
    plugins = attacks_plugins()
    if resultJson:
        resultJson = {}
    startup_plugins(plugins)
    REPORT['attacks'] = resultJson

def startup_spec_audit(plugin:str):
    if plugin in audit_info.keys():
        plugins = spec_audit_plugins(plugin)
        startup_plugins(plugins)
    else:
        warn("模块不存在!")

def startup_full_audit():
    global resultJson
    if not REPORT['startTime']:
        REPORT['startTime'] = strftime("%Y/%m/%d at %H:%M:%S")
    execmod.append("audit")
    plugins = audit_plugins()
    if resultJson:
        resultJson = {}
    startup_plugins(plugins)
    REPORT['audit'] = resultJson

def startup_spec_brute(p:str):
    if p in brute_info.keys():
        plugins = spec_brute_plugins(p)
        startup_plugins(plugins)
    else:
        warn("模块不存在!")

def startup_full_brute():
    global resultJson
    if not REPORT['startTime']:
        REPORT['startTime'] = strftime("%Y/%m/%d at %H:%M:%S")
    execmod.append("brute")
    plugins = brute_plugins()
    if resultJson:
        resultJson = {}
    startup_plugins(plugins)
    REPORT['brute'] = resultJson

def startup_spec_disclosure(p:str):
    if p in disclosure_info.keys():
        plugins = spec_disclosure_plugins(p)
        startup_plugins(plugins)
    else:
        warn("模块不存在!")

def startup_full_disclosure():
    global resultJson
    if not REPORT['startTime']:
        REPORT['startTime'] = strftime("%Y/%m/%d at %H:%M:%S")
    execmod.append("disclosure")
    plugins = disclosure_plugins()
    if resultJson:
        resultJson = {}
    startup_plugins(plugins)
    REPORT['disclosure'] = resultJson

def get_argv():
    for k,v in ARGV.items():
        s = "{}:{}".format(k,str(v))
        plus(s)
def is_set_argv(key):
    if key in ARGV.keys():
        if ARGV[key]:
            return True
    else:
        return False

def write_json(jsonData):
    name = strftime('%m%d%H%M%S')
    filepath = REPORT_JSON + sep + name +".json"
    with open(filepath,'w',encoding="utf-8") as f:
        json.dump(jsonData,f)
    
def generate_report(filename,title="安全扫描报告"):
    global execmod
    if filename and REPORT['startTime'] and execmod:
        REPORT['fileName'] = filename
        REPORT['name'] = title
        REPORT['title'] = title
        execmod = list(set(execmod))
        REPORT['execmod'] = execmod
        total = success = fail = 0
        for i in execmod:
            if REPORT[i]:
                for k,v in REPORT[i].items():
                    if v:
                        success += 1
                    else:
                        fail += 1
        total = success + fail
        REPORT['total'] = str(total)
        REPORT['successTotal'] = str(success)
        REPORT['failTotal'] = str(fail)
        report_gen = HtmlReport(REPORT)
        report_gen.generate_report()
    else:
        warn("生成报告失败!")

        



